import socket
import asyncio
from typing import Callable
from pathlib import Path
import sys
BASE_DIR = Path(__file__).resolve().parent.parent.parent
sys.path.append(str(BASE_DIR).replace("\\", "\\\\"))
from agent.pkt import Pkt, Cmd
from agent.session import Session
from agent.session_admin import SessionAdmin
from agent.kcp_py.kcp.server import Connection,KCPServerAsync

sa = SessionAdmin()

class AgentServer(object):
  read_buf_max = 1024
  timeout_max = 5
  global_id = 0

  def __init__(self):
    pass

  @staticmethod
  def get_add_global_id()->int:
    AgentServer.global_id = AgentServer.global_id + 1
    if AgentServer.global_id == 0:
      AgentServer.global_id = 1

    return AgentServer.global_id

  async def on_start(self):
    pass

  @staticmethod
  def getHostAndPort(data):
    is_connect = False
    if data.find(b'CONNECT') >= 0:
      is_connect = True
    i0 = data.find(b'Host')
    if i0 <= 0:
      return False
    data = data[i0+6:]
    i1 = data.find(b'\r\n')
    data = data[:i1]
    i2 = data.find(b':')
    if i2 > 0:
      return True, data[:i2], 443, is_connect
    return True, data, 80, is_connect
    
  async def on_data(self, connection: Connection, data: bytes) -> None:
    print(f"Received data from {connection.address}")
    
    pkt = Pkt(0, 0, 0)
    if pkt.unpack(data) is False:
      print('unpack failed')
      return

    if pkt.cmd == Cmd.QUERY_GLOBALID:
      pkt.globalid = AgentServer.get_add_global_id()
      pkt.data = None
      connection.enqueue(data=pkt.pack())
      return
    if pkt.cmd == Cmd.KEEP_ALIVE:
      print('recv keepalive from global id ', pkt.globalid)
      pkt.data = None
      connection.enqueue(data=pkt.pack())
      return
    s = sa.get_session_by_ids(pkt.globalid, pkt.localid)
    if s:
      s.writer.write(pkt.data)
      return

    #转发 
    try:
      valid, host, port, is_connect = AgentServer.getHostAndPort(pkt.data)
    except BaseException as e:
      print(e)
      return
    if valid is False:
      print('err data')
      return
    try:
      reader, writer = await asyncio.open_connection(host, port)
    except BaseException as e:
      print(e)
      return
    if reader != None:
      s = Session(pkt.globalid, pkt.localid, reader, writer)
      sa.add_session(s)
      if is_connect:
        pkt.data = b'HTTP/1.1 200 Connection Established\r\nConnection: close\r\n\r\n'
        connection.enqueue(pkt.pack())
      else:
        writer.write(pkt.data)
      
      while True:
        try:
          data = await reader.read(AgentServer.read_buf_max)
          if len(data) <= 0:
            break
          pkt.data = data
          connection.enqueue(pkt.pack())
        except:
          break
      writer.close()
      await asyncio.sleep(5)
      sa.del_session(s)


server = KCPServerAsync(
    "0.0.0.0",
    5100,
    conv_id=1,
)

agent = AgentServer()

@server.on_start
async def on_start() -> None:
    print("Server started!")
    await agent.on_start()

@server.on_data
async def on_data(connection: Connection, data: bytes) -> None:
  loop = asyncio.get_running_loop()
  loop.create_task(agent.on_data(connection, data))

server.start()











